from classifier.classifier import Classifier
from sklearn.naive_bayes import GaussianNB
from sklearn.naive_bayes import BernoulliNB
from sklearn.naive_bayes import MultinomialNB


# 朴素贝叶斯
class NaiveBayes(Classifier):
    # 模型类型(GaussianNB,BernoulliNB,MultinomialNB)
    model_type = "GaussianNB"
    
    def __init__(self):
        Classifier.__init__(self)
        self.algorithm_name = "朴素贝叶斯"
        self.ipynb_template_name = "naive_bayes-template.ipynb"
        
    def implent(self): 
        Classifier.implent(self)
        # 构造模型
        self.algorithm = None
        if self.model_type == "GaussianNB":
            self.algorithm = GaussianNB()
        elif self.model_type == "BernoulliNB":
            self.algorithm = BernoulliNB()
        else:
            self.algorithm = MultinomialNB()
        # 训练模型
        self.algorithm.fit(self.train_inputs, self.train_outputs)
        # 评估模型
        self.score = self.algorithm.score(self.test_inputs, self.test_outputs)
        # 预测
        self.predict_output_values = self.algorithm.predict(self.predict_input_values)
        self.predict_output_probas = self.algorithm.predict_proba(self.predict_input_values)
    
    def prepareIpynbItems(self):
        Classifier.prepareIpynbItems(self)
        self.ipynb_items["#model_type#"] = self.model_type
